package fem2;

import inf.jlinalg.IMatrix;

import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import fem2.enu.EchoLevelType;
import fem2.enu.MassMatrixType;

/**
 * The Model class for nonlinear analysis
 * 
 * @author hbui
 * 
 */

public class Model extends Echoer implements ElementPart {

	private Mesh mesh;
	private Assembler as;
	private ArrayList<Element> elementList = new ArrayList<Element>();
	private ArrayList<Load> loadList = new ArrayList<Load>();

	private Set<String> elemTypes = new HashSet<String>();
	private Map<String, Integer> elemNumNodes = new HashMap<String, Integer>();

	/**
	 * @param m
	 */
	public Model(Mesh m) {
		this.mesh = m;
		this.as = new SequentialAssembler(this);
	}

	/**
	 * @return Gets the mesh.
	 */
	public Mesh getMesh() {
		return mesh;
	}

	/**
	 * Adds an element
	 */
	public Element addElement(Element e) {
		elementList.add(e);
		e.setId(elementList.size());
		elemTypes.add(e.getElemType());
		elemNumNodes.put(e.getElemType(), e.getMeshPart().countModes());
		return e;
	}

	/**
	 * Returns the number of elements
	 * 
	 * @return number of elements.
	 */
	public int countElements() {
		return elementList.size();
	}

	/**
	 * Gets a an element from the Model.
	 * 
	 * @param i
	 *            the index of element.
	 * @return the element corresponding to given index.
	 */
	public Element getElement(int i) {
		return elementList.get(i);
	}

	/**
	 * 
	 * @return set of element types in the model
	 */
	public Set<String> getElementTypes() {
		return this.elemTypes;
	}

	/**
	 * 
	 * @return map from element type to number of nodes on that element
	 */
	public Map<String, Integer> getElementNumNodes() {
		return this.elemNumNodes;
	}

	/**
	 * Adds the load to the model.
	 * 
	 * @param l
	 *            the load needs to be added.
	 * @return added load.
	 */
	public void addLoad(Load l) {
		this.loadList.add(l);
		l.setId(loadList.size());
	}

	/**
	 * Returns the number of loads in the mesh
	 * 
	 * @return number of loads.
	 */
	public int countLoads() {
		return this.loadList.size();
	}

	/**
	 * Gets a load from mesh.
	 * 
	 * @param i
	 *            the index of mesh part.
	 * @return the load corresponding to given index.
	 */
	public Load getLoad(int i) {
		return this.loadList.get(i);
	}

	/**
	 * enumerates the degrees of freedom
	 */
	public void preAnalysis() {
		this.mesh.enumerateDofs();
	}

	/**
	 * 
	 */
	public void postAnalysis() {
		/*
		 * TODO
		 */
	}

	/**
	 * assemble r = r + r_model. User are responsible to clear the external
	 * force vector before assembling
	 * 
	 * @param r
	 *            global load vector
	 * @param t
	 *            time
	 * @return pointer to external load vector
	 */
	public double[] assembleR(double[] r, double t) {
		for (int k = 0; k < this.countLoads(); k++) {
			Load load = this.getLoad(k);

			if (!load.isInactive()) {
				double[] re = load.computeRe(t);

				int[] dofs = load.getMeshPart().getDofNumbers();
				// for (int j = 0; j < dofs.length; j++) {
				// if (dofs[j] != -1) {
				// r[dofs[j]] += re[j];
				// }
				// }
				/*
				 * this is to account for Xfem
				 */
				for (int j = 0; j < re.length; j++) {
					if (dofs[j] != -1) {
						r[dofs[j]] += re[j];
					}
				}
			}

		}
		return r;
	}

	/**
	 * assemble r = r + alpha*r_model. User is responsible to clear the external
	 * force vector before assembling
	 * 
	 * @param alpha
	 * @param r
	 *            global load vector
	 * @param t
	 *            time
	 * @return pointer to external load vector
	 */
	public double[] assembleR(double alpha, double[] r, double t) {
		for (int k = 0; k < this.countLoads(); k++) {
			Load load = this.getLoad(k);

			if (!load.isInactive()) {
				double[] re = load.computeRe(t);

				int[] dofs = load.getMeshPart().getDofNumbers();
				for (int j = 0; j < dofs.length; j++) {
					if (dofs[j] != -1) {
						r[dofs[j]] += alpha * re[j];
					}
				}
			}

		}
		return r;
	}

	/**
	 * Assembles the internal force vector.
	 * 
	 * @param ri
	 *            : global internal force vector
	 * @return pointer to internal load vector
	 */
	public double[] assembleRi(double[] ri) {
		long start = System.currentTimeMillis();
		as.assembleRi(ri);
		long end = System.currentTimeMillis();
		if (getEchoLevel().contain(EchoLevelType.AUX1)) {
			System.out.println("Assemble Ri time: " + (end - start) + " ms.");
		}
		return ri;
	}

	/**
	 * Assembles the element stiffness matrices into the global stiffness
	 * matrix.
	 * 
	 * @param kk
	 *            global stiffness matrix, must contain only zeroes on entry.
	 */
	public void assembleKt(IMatrix kk) {
		long start = System.currentTimeMillis();
		as.assembleKt(kk);
		long end = System.currentTimeMillis();
		if (getEchoLevel().contain(EchoLevelType.AUX1)) {
			System.out.println("Assemble Kt time: " + (end - start) + " ms.");
		}
	}

	public void assembleM(IMatrix m, MassMatrixType massType) {
		long start = System.currentTimeMillis();
		as.assembleM(m, massType);
		long end = System.currentTimeMillis();
		if (getEchoLevel().contain(EchoLevelType.AUX1)) {
			System.out.println("Assemble M time: " + (end - start) + " ms.");
		}
	}

	/**
	 * save this model to an object file
	 * 
	 * TODO: doesn't work yet
	 * 
	 * @param str
	 *            file name
	 * @throws IOException
	 */
	public void Save(String str) throws IOException {
		FileOutputStream fileOut = new FileOutputStream(str);
		ObjectOutputStream out = new ObjectOutputStream(fileOut);
		out.writeObject(this);
		out.flush();
		out.close();
		fileOut.close();
	}

	/**
	 * 
	 * @return the size of linear equation system
	 */
	public int getSize() {
		return mesh.countDofs();
	}

	/**
	 * apply the solution vector to the system at time t
	 * 
	 * @param t
	 *            time
	 * @param x
	 *            solution vector
	 */
	public void setX(double t, double[] x) {
		mesh.setUHatGlobal(t, x);
	}

	/**
	 * set assembler for this model
	 * 
	 * @param as
	 */
	public void setAssembler(Assembler as) {
		this.as = as;
	}

	/************************************
	 * ElementPart Override
	 ************************************/

	@Override
	public void Initialize() {
		preAnalysis();
		for (Element e : elementList) {
			e.Initialize();
		}
	}

	@Override
	public void InitializeSolutionStep() {
		for (Element e : elementList) {
			e.InitializeSolutionStep();
		}
	}

	@Override
	public void InitalizeNonlinearIteration() {
		for (Element e : elementList) {
			e.InitalizeNonlinearIteration();
		}
	}

	@Override
	public void FinalizeNonlinearIteration() {
		for (Element e : elementList) {
			e.FinalizeNonlinearIteration();
		}
	}

	public void FinalizeSolutionStep() {
		for (Element e : elementList) {
			e.FinalizeSolutionStep();
		}
		postAnalysis();
	}

	@Override
	public void CalculateOutputData() {
	}

}
